# Figure 1, Tables A2, A3

suppressPackageStartupMessages({
  library(tidyverse)
  library(haven)
  library(urbnmapr) # remotes::install_github("UrbanInstitute/urbnmapr")
  library(sf)
  library(gt)
})

cat_lbl = c("state partisan" = "State Partisan Offices",
            "state nonpartisan" = "State Nonpartisan Offices",
            "state measure" = "State Ballot Measures",
            "local partisan" = "Local Partisan Offices",
            "local nonpartisan" = "Local Nonpartisan Offices",
            "local measure" = "Local Ballot Measures")

## counts ----
## number of counties
state_counties <- read_dta("APSR_counties_used.dta") |> 
  summarize(counties = n(), .by = state) 

## number of voters
state_voters <- read_dta("tbA2_nvoters_by-county.dta") |> 
  filter(office_type != "natl_partisan_races") |> 
  summarize(voters = sum(voters), .by = c(office_type, state)) |> 
  mutate(office_type = word(str_replace_all(office_type, "_", " "), 1, 2),
         office_type = str_replace(office_type, "ballot", "measure")) |> 
  filter(voters > 100)

## number of contests
contests_long <- c(   
  `local measure` = "wide_local_ballot_measures.dta",
  `local nonpartisan` = "wide_local_nonpartisan_races.dta",
  `local partisan` = "wide_local_partisan_races.dta",
  `state measure` = "wide_state_ballot_measures.dta",
  `state nonpartisan` = "wide_state_nonpartisan_races.dta",
  `state partisan` = "wide_state_partisan_races.dta"
) |> 
  map(read_dta) |> 
  list_rbind(names_to = "office_type")

state_contests_all <- contests_long |>   
  summarize(
    contests = n(),
    .by = c(office_type, state)
  )

### compet
state_con_compet <- contests_long |> 
  mutate(compet = coalesce(
    Y_choice / (Y_choice + N_choice),
    A_choice / (A_choice + B_choice),
    D_choice / (D_choice + R_choice)
  )) |> 
  tidylog::filter(between(compet, 0.05, 0.95)) |> 
  summarize(
    contests_compet = n(),
    .by = c(office_type, state)
  )

state_contests <- state_contests_all |> 
  left_join(state_con_compet) |> 
  mutate(contests_compet = replace_na(contests_compet, 0)) |> 
  mutate(contests_frac = glue::glue("{contests_compet} / {contests}"))

## Number of voters
# deduplicate, usualy pick state partisan but pick ballot measure for MD and NJ
v_partition <-  read_dta("tbA2_nvoters_by-county.dta") |> 
  filter(office_type != "natl_partisan_races") |> 
  slice_max(voters, by = c(state, county), with_ties = FALSE)

N_voters <- v_partition |> 
  pull(voters) |> sum()
N_voters |> 
  write_lines(file.path(PATH_FIG, "../numbers/votersN_pres.tex"))

round(N_voters/1e6, 1)
v_partition |> select(matches("pid")) |> colSums() |> prop.table()


# map ------

# shapefiles
state_shp <- get_urbn_map("states", sf = TRUE) |> 
  rename(state = state_abbv)
state_sf <- state_shp |> 
  left_join(state_voters, by = c("state"), relationship = "one-to-many") |> 
  complete(nesting(state, geometry), office_type, fill = list(voters = NA_real_)) |> 
  mutate(office_type = factor(office_type, levels = names(cat_lbl))) |> 
  filter(!is.na(office_type)) |> 
  st_sf()

## Labelling
tot_voters <- state_voters |> 
  summarize(voters = sum(voters), .by = office_type) 
cat_lbl2 <- enframe(cat_lbl, name = "office_type", value = "lab") |> 
  left_join(tot_voters) |> 
  mutate(lbl = glue::glue("**{lab}**<br>({scales::comma(voters)} voters)")) |> 
  select(office_type, lbl) |> 
  deframe()

gg_st <- state_sf |> 
  ggplot(aes(fill = voters)) +
  geom_sf() +
  facet_wrap(~ office_type, labeller = labeller(office_type = cat_lbl2)) +
  scale_fill_fermenter(
    breaks = c(1e3, 1e4, 1e5, 1e6, 2e6, 4e6, 8e6), na.value = "white", 
    palette = "Oranges",
    labels = scales::comma,
    direction = 1) +
  theme_void() +
  theme(strip.text = ggtext::element_markdown()) +
  labs(fill = "Voters with\nNational\nPartisanship")


# tables -----
state_key <- state.name; names(state_key) <- state.abb
bystate_gt <- function(tbl, var, counties = state_counties) {
  wide <- tbl |> 
    filter(office_type != "natl partisan") |> 
    left_join(counties, relationship = "many-to-one", by = "state") |> 
    mutate(office_type = factor(office_type, levels = names(cat_lbl))) |> 
    arrange(office_type) |> 
    filter(!is.na(office_type)) |>  
    pivot_wider(id_cols = c(state, counties), 
                names_from = office_type, 
                values_from = {{var}}) |> 
    relocate(state, counties, matches(" partisan"), matches("nonpartisan")) |> 
    arrange(state) |> 
    mutate(state = state_key[state])
  
  # if var is a variable called contests_frac
  # var <- quote(var)
  if (rlang::as_label(enquo(var)) == "contests_frac") {
    summ_stats <- state_contests |> 
      summarize(
        .by = office_type,
        contests_frac = glue::glue("{scales::comma(sum(contests_compet))} / {scales::comma(sum(contests))}")) |> 
      mutate(counties = sum(state_counties$counties),
             state = "Total")
  } 
  
  if (rlang::as_label(enquo(var)) == "voters") {
    summ_stats <- state_voters |> 
      summarize(
        .by = office_type,
        voters = sum(voters)) |> 
      mutate(counties = sum(state_counties$counties),
             state = "Total")
  }
  
  ## stack
  wide_summ <- summ_stats |> 
    pivot_wider(id_cols = c(state, counties), names_from = office_type, values_from = {{var}})
  
  wide <- wide |> bind_rows(wide_summ)
  
  ## table
  wide_gt <- wide |>  
    gt() |> 
    sub_missing(missing_text = "-") |> 
    fmt_integer() |> 
    cols_label(
      starts_with("counties") ~ "Ctys.",
      starts_with("state") ~ "State",
      starts_with("local") ~ "Local",
    ) |> 
    tab_spanner("Partisan Offices", matches("( partisan|federal)")) |> 
    tab_spanner("Non-partisan Offices", ends_with("nonpartisan")) |> 
    tab_spanner("Ballot Measures", ends_with("measure")) |> 
    tab_style(
      style = cell_borders(
        sides = "bottom"
      ),
      locations = cells_body(rows = nrow(wide) - 1)
    )
  
  wide_gt
}


# Export -----
# Uncomment the `gtsave` line, fix the path to your local, and 
# Paste only the contents (not longtable header) to paper

bystate_gt(state_voters, voters) 
# gtsave(file.path("~/Downloads", "tables/table_A2_counts_by-state_voters.tex"))

bystate_gt(state_contests, contests_frac) 
# gtsave(file.path("~/Downloads", "tables/table_A3_counts_by-state_contests.tex"))

gg_st
# ggsave(file.path(PATH_FIG, "fg1_map_stateN.pdf"), gg_st, w = 7.5, h = 4)
